#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
WebSocket代理桥接服务快速启动和验证脚本
用于快速验证所有功能是否正常工作
"""

import asyncio
import subprocess
import time
import logging
import sys
import os
import json
import aiohttp
from datetime import datetime

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class QuickStartValidator:
    """快速启动验证器"""
    
    def __init__(self):
        self.processes = []
        self.api_url = "http://127.0.0.1:8998"
    
    def cleanup(self):
        """清理所有进程"""
        logger.info("清理进程...")
        for process in self.processes:
            try:
                if process.poll() is None:
                    process.terminate()
                    process.wait(timeout=3)
            except subprocess.TimeoutExpired:
                process.kill()
            except Exception as e:
                logger.error(f"清理进程时出错: {str(e)}")
        self.processes.clear()
    
    def start_process(self, command: list, name: str, wait_time: int = 2):
        """启动进程"""
        try:
            logger.info(f"启动 {name}...")
            process = subprocess.Popen(
                command,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True
            )
            self.processes.append(process)
            time.sleep(wait_time)
            
            if process.poll() is None:
                logger.info(f"✓ {name} 启动成功")
                return True
            else:
                stdout, stderr = process.communicate()
                logger.error(f"✗ {name} 启动失败")
                if stderr:
                    logger.error(f"错误信息: {stderr}")
                return False
        except Exception as e:
            logger.error(f"启动 {name} 时出错: {str(e)}")
            return False
    
    async def check_api_health(self):
        """检查API健康状态"""
        try:
            async with aiohttp.ClientSession() as session:
                async with session.get(f"{self.api_url}/health", timeout=5) as response:
                    if response.status == 200:
                        data = await response.json()
                        logger.info("✓ HTTP API服务正常")
                        logger.info(f"  活跃连接数: {data.get('active_connections', 0)}")
                        return True
                    else:
                        logger.error(f"✗ HTTP API响应异常: {response.status}")
                        return False
        except Exception as e:
            logger.error(f"✗ HTTP API检查失败: {str(e)}")
            return False
    
    async def validate_core_functionality(self):
        """验证核心功能"""
        logger.info("=" * 60)
        logger.info("验证核心功能")
        logger.info("=" * 60)
        
        # 1. 检查依赖
        logger.info("1. 检查Python依赖...")
        try:
            import mitmproxy
            import aiohttp
            logger.info("✓ 所有依赖已安装")
        except ImportError as e:
            logger.error(f"✗ 缺少依赖: {str(e)}")
            logger.error("请运行: pip install -r requirements.txt")
            return False
        
        # 2. 检查文件完整性
        logger.info("2. 检查文件完整性...")
        required_files = [
            'websocket_proxy_bridge.py',
            'test_websocket_server.py',
            'test_websocket_client.py',
            'test_http_api.py',
            'example_usage.py',
            'requirements.txt',
            'README.md'
        ]
        
        missing_files = []
        for file in required_files:
            if not os.path.exists(file):
                missing_files.append(file)
        
        if missing_files:
            logger.error(f"✗ 缺少文件: {', '.join(missing_files)}")
            return False
        else:
            logger.info("✓ 所有必需文件存在")
        
        # 3. 启动WebSocket测试服务器
        logger.info("3. 启动WebSocket测试服务器...")
        server_started = self.start_process(
            ['python', 'test_websocket_server.py', '--host', '127.0.0.1', '--port', '8765', '--no-periodic'],
            'WebSocket测试服务器',
            wait_time=3
        )
        
        if not server_started:
            logger.error("✗ WebSocket测试服务器启动失败")
            return False
        
        # 4. 启动代理桥接服务
        logger.info("4. 启动WebSocket代理桥接服务...")
        proxy_started = self.start_process(
            ['python', 'websocket_proxy_bridge.py', '--proxy-port', '8999', '--api-port', '8998', '--log-level', 'WARNING'],
            'WebSocket代理桥接服务',
            wait_time=5
        )
        
        if not proxy_started:
            logger.error("✗ WebSocket代理桥接服务启动失败")
            return False
        
        # 5. 等待服务完全启动
        logger.info("5. 等待服务完全启动...")
        await asyncio.sleep(3)
        
        # 6. 验证HTTP API
        logger.info("6. 验证HTTP API...")
        api_healthy = await self.check_api_health()
        
        if not api_healthy:
            logger.error("✗ HTTP API验证失败")
            return False
        
        # 7. 测试API端点
        logger.info("7. 测试API端点...")
        try:
            async with aiohttp.ClientSession() as session:
                # 测试连接列表端点
                async with session.get(f"{self.api_url}/connections") as response:
                    if response.status == 200:
                        data = await response.json()
                        logger.info(f"✓ 连接列表API正常，当前连接数: {data.get('count', 0)}")
                    else:
                        logger.error(f"✗ 连接列表API异常: {response.status}")
                        return False
                
                # 测试发送消息端点（预期失败，因为没有连接）
                test_payload = {
                    "connection_id": "test-id",
                    "message": "test message"
                }
                async with session.post(f"{self.api_url}/send", json=test_payload) as response:
                    if response.status == 404:  # 预期的错误，因为连接不存在
                        logger.info("✓ 发送消息API正常（正确返回404错误）")
                    else:
                        logger.warning(f"发送消息API返回意外状态: {response.status}")
                
        except Exception as e:
            logger.error(f"✗ API端点测试失败: {str(e)}")
            return False
        
        logger.info("=" * 60)
        logger.info("✓ 所有核心功能验证通过！")
        logger.info("=" * 60)
        
        return True
    
    async def show_usage_instructions(self):
        """显示使用说明"""
        logger.info("=" * 60)
        logger.info("使用说明")
        logger.info("=" * 60)
        
        logger.info("服务已启动，您可以:")
        logger.info("")
        logger.info("1. 配置WebSocket客户端使用代理:")
        logger.info("   - HTTP代理: 127.0.0.1:8999")
        logger.info("   - HTTPS代理: 127.0.0.1:8999")
        logger.info("")
        logger.info("2. 使用HTTP API控制WebSocket连接:")
        logger.info("   - 健康检查: curl http://127.0.0.1:8998/health")
        logger.info("   - 查看连接: curl http://127.0.0.1:8998/connections")
        logger.info("   - 发送消息: curl -X POST http://127.0.0.1:8998/send \\")
        logger.info("                    -H 'Content-Type: application/json' \\")
        logger.info("                    -d '{\"connection_id\":\"ID\",\"message\":\"Hello!\"}'")
        logger.info("")
        logger.info("3. 运行示例程序:")
        logger.info("   - python example_usage.py --mode demo")
        logger.info("   - python example_usage.py --mode interactive")
        logger.info("")
        logger.info("4. 运行测试工具:")
        logger.info("   - python test_http_api.py")
        logger.info("   - python test_websocket_client.py --url ws://127.0.0.1:8765")
        logger.info("")
        logger.info("按 Ctrl+C 停止所有服务")
        logger.info("=" * 60)
    
    async def run_quick_start(self):
        """运行快速启动流程"""
        try:
            # 验证核心功能
            validation_success = await self.validate_core_functionality()
            
            if validation_success:
                # 显示使用说明
                await self.show_usage_instructions()
                
                # 保持服务运行
                try:
                    while True:
                        await asyncio.sleep(1)
                except KeyboardInterrupt:
                    logger.info("收到中断信号，正在关闭服务...")
            else:
                logger.error("验证失败，请检查错误信息并修复问题")
                return False
                
        except Exception as e:
            logger.error(f"快速启动过程中出错: {str(e)}")
            return False
        finally:
            self.cleanup()
        
        return validation_success


async def main():
    """主函数"""
    logger.info("WebSocket代理桥接服务 - 快速启动验证")
    logger.info(f"启动时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    
    validator = QuickStartValidator()
    
    try:
        success = await validator.run_quick_start()
        exit_code = 0 if success else 1
        logger.info(f"快速启动完成，退出码: {exit_code}")
        sys.exit(exit_code)
    except KeyboardInterrupt:
        logger.info("用户中断")
        validator.cleanup()
        sys.exit(0)
    except Exception as e:
        logger.error(f"快速启动异常: {str(e)}")
        validator.cleanup()
        sys.exit(1)


if __name__ == "__main__":
    asyncio.run(main())